'''
Created on Apr 6, 2012

@author: Rafael
'''
import os
import time
from nipy import load_image
from nipy import save_image
from scipy.cluster.vq import kmeans2
import numpy as np
import scipy as sp
from nipy.core.api import Image, vox2mni

bd = './Inputs'
os.chdir(bd)

inmask='Sub_1_GM.nii.gz'

in_mask = load_image(inmask)

mask=in_mask.get_data()
shape=mask.shape

print "Mask dimensions: ", mask.shape

voxels=shape[0]*shape[1]*shape[2]
print 'original voxels: ',  voxels


def GetCluster(Cluster,mask, shape, voxels):
    i = 0
    MS_Vector=np.zeros( ( voxels) )
    while i < voxels-1:
        for z in range(shape[2]):
            for y in range(shape[1]):
                for x in range(shape[0]):
                    MS_Vector[i]=mask[x,y,z]
                    i+=1
    
    CL=np.zeros(shape[0], shape[1], shape[2])
    i = 0
    j = 0
    while i < voxels-1:
        for z in range(shape[2]):
            for y in range(shape[1]):
                for x in range(shape[0]):
                    if MS_Vector[i] != 0:
                        CL[x,y,z]=Cluster[j]
                        j += 1
    return CL


Corr_Matrix=np.load('correlation_matrix')
print "Corr_Matrix shape", Corr_Matrix.shape

##################################
nclusters=4

print "begin clusters at: ", time.ctime()
clusters=kmeans2(Corr_Matrix, nclusters)
np.save('4-clusters', clusters)
print "end 4 clusters, 10 iterations at: ", time.ctime()

#################################

cluster1=GetCluster(clusters[0:],mask, shape, voxels)
arr_img = Image(cluster1, vox2mni(np.eye(4)))
saved_cluster1= save_image(arr_img, 'cluster1')

